"""
https://www.tensorflow.org/guide/keras/train_and_evaluate
"""

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers, losses, metrics, optimizers, Model
from python_ai.common.xcommon import *
import os
import sys
from sklearn.datasets import load_breast_cancer
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

# data
x, y = load_breast_cancer(return_X_y=True)
x = StandardScaler().fit_transform(x)

x_train, x_test, y_train, y_test = train_test_split(x, y, train_size=0.7, random_state=777)
m, n = x_train.shape

# Reserve 10,000 samples for validation
n_val = int(np.ceil(m * 0.1))
x_val = x_train[-n_val:]
y_val = y_train[-n_val:]
x_train = x_train[:-n_val]
y_train = y_train[:-n_val]

# model
inputs = keras.Input(shape=(n,), name="digits")
x = layers.Dense(64, activation="relu", name="dense_1")(inputs)
x = layers.Dense(64, activation="relu", name="dense_2")(x)
outputs = layers.Dense(1, activation="sigmoid", name="predictions")(x)

model = Model(inputs=inputs, outputs=outputs)

# specify the training config
model.compile(
    optimizer=optimizers.RMSprop(),
    loss=losses.binary_crossentropy,
    metrics=[metrics.binary_accuracy]
)

# fit and train
sep('Fit and train')
ver = 'v1.0'
filename = os.path.basename(__file__)
log_dir = os.path.join('_log', filename, ver)
os.makedirs(log_dir, exist_ok=True)
tb_callback = tf.keras.callbacks.TensorBoard(
    log_dir=log_dir,
    update_freq='batch',
    profile_batch=0
)
early_stop = tf.keras.callbacks.EarlyStopping(
    monitor='val_binary_accuracy',
    min_delta=1e-4,
    patience=2,
    verbose=1
)
history = model.fit(
    x_train,
    y_train,
    batch_size=64,
    epochs=20,
    validation_data=(x_val, y_val),
    callbacks=[tb_callback, early_stop]
)
print(history.history)

sep('Evaluate on test data')
results = model.evaluate(x_test, y_test, batch_size=128)
print(f'test loss, test acc: {results}')

